# -*- coding: utf-8 -*-
import math
import re
import sys
import numpy as np
import pandas as pd
import pyqtgraph as pg
from PyQt5.QtWidgets import QVBoxLayout, QWidget, QApplication
from src.sharedSrc import rHEX, getSetting
from pyqtgraph.exporters import ImageExporter


class WaferSummary(QWidget):
    def __init__(self):
        super().__init__()

        self.text_items = []
        self.plotlayout = pg.GraphicsLayoutWidget()
        self.plotlayout.setBackground('white')  # Set the entire plot area background to white
        self.setLayout(QVBoxLayout())
        self.layout().addWidget(self.plotlayout)
        self.view_list = []
        self.settings = getSetting()

    def setSummaryData(self, binData: dict, filename: str = ''):
        if not binData:
            return
        row = 0
        for binType in ["ND", "RD", 'ALLND', 'ALLRD']:
            if binType not in binData:
                continue
            hRD = binData[binType]  # hRD {0: {1: 26, 3: 61, 4: 57, 5: 286, 9: 15}}
            if len(hRD[1]) == 0:
                continue
            binTicks = binData[
                binType + "_Ticks"]  # hRD {0: {1: 26, 3: 50, 9: 15, 200: 1, 201: 2, 202: 8, 300: 2, 310: 13, 311: 1,
            # 320: 2, 330: 12, 340: 2, 350: 1, 360: 4, 362: 1, 370: 1, 410: 18, 500: 39, 700: 5, 800: 18, 1000: 2,
            # 2000: 150, 3000: 55, 4000: 8, 6000: 5, 7000: 2, 8000: 2}}
            isND = True if "ND" in binType else False

            num_files = len(hRD)
            tmpVbList = []
            # binColorDict = settings.NDColor if isND else settings.RDColor
            binTypeName = "Normal Distribution" if isND else "Resistance Distribution"
            self.plotlayout.addLabel(f"{binTypeName}",
                                     row=row, col=1,
                                     rowspan=1, colspan=num_files,
                                     size="20pt")
            row += 1
            for fid in hRD.keys():
                view_bin = pg.ViewBox()
                # view_bin.invertY(True)
                pitem = pg.PlotItem(viewBox=view_bin)
                binStats = hRD[fid]
                if len(binStats) == 0:
                    continue
                numList = binTicks.keys()
                cntList = np.array([binStats.get(n, 0) for n in numList])
                x = np.arange(len(numList))
                if isND:
                    colorList = []
                    self.Color = self.settings.wafer_dict
                    for i in numList:
                        if i in self.Color:
                            colorList.append(self.Color[i])
                        else:
                            colorList.append(rHEX())
                    width = 0.08  # Adjust width to your preference
                    bar = pg.BarGraphItem(x=x, height=cntList, width=width, brushes=colorList)
                    pitem.addItem(bar)
                    # # Add text items above bars
                    for index, value in enumerate(cntList):
                        if value == 0:
                            continue
                        text_item = pg.TextItem(text=str(value), color='black', anchor=(0.5, 0.5))
                        text_item.setPos(x[index], value + 0.25)
                        pitem.addItem(text_item)
                        # Format ticks with labels from binTicks

                    ticks = [[binTicks[n] for n in numList]]
                    pitem.getAxis("bottom").setTicks(ticks)
                else:
                    # Plot the normal distribution curve
                    line_plot = pg.PlotCurveItem(x, cntList, pen=pg.mkPen(color='r', width=2))
                    data = list(binTicks.values())
                    pitem.addItem(line_plot)
                    min_value = binData[binType + 'min']
                    max_value = binData[binType + 'max']
                    testName = binData[binType + 'testName']
                    if abs(min_value) > 1:
                        pitem.getAxis("bottom").setLabel(f'{testName}: [{min_value:.3f}~{max_value:.3f}] Unit: (mA/mV)')
                    else:
                        pitem.getAxis("bottom").setLabel(f'{testName}: [{min_value:.6f}~{max_value:.6f}] Unit: (mA/mV)')
                    len_data = len(binTicks)
                    len_num = math.floor(math.log10(len_data)) + 1  #  计算数据的位数
                    space = 10 if len_num <= 1 else 10**(len_num-1)
                    scatter_size = 10-len_num if 10-len_num else 1
                    point_plot = pg.ScatterPlotItem(x, cntList, pen=pg.mkPen(None), brush=pg.mkBrush(255, 0, 0),
                                                    size=scatter_size, hoverable=True, data=data)
                    pitem.addItem(point_plot)
                    # Optimized tick calculation using list comprehension
                    # ticks = [binTicks[i] for i in numList if i % space == 0]
                    ticks = []
                    distinct_ticks = {}
                    for i in numList:
                        if i % space == 0 and binTicks[i][1] not in distinct_ticks:
                            distinct_ticks[binTicks[i][1]] = 0
                            ticks.append(binTicks[i])
                    # ticks.append(list(distinct_ticks.values()))
                    # ticks.append(binTicks[len_data-1])
                    # print(ticks)
                    # Set ticks in one call outside of the loop
                    pitem.getAxis("bottom").setTicks([ticks])

                x_max = len(numList)
                y_max = max(cntList) * 1.15
                view_bin.setLimits(xMin=-1, xMax=x_max,
                                   yMin=0, yMax=y_max,
                                   minXRange=4, minYRange=2)
                view_bin.setRange(xRange=(-1, x_max),
                                  yRange=(0, y_max),
                                  disableAutoRange=False)
                self.plotlayout.addItem(pitem, row=row, col=fid, rowspan=1, colspan=1)
                tmpVbList.append(view_bin)
                self.view_list.append(view_bin)
            row += 1

        if filename != '':
            self.save_plot(filename)

    def save_plot(self, filename: str) -> None:
        exporter = ImageExporter(self.plotlayout.scene())
        exporter.export(filename)


def wafer_data_summary(sample_wafer_data: dict) -> dict:
    ND = {}
    ND_Ticks = {}
    RD = {}
    RD_Ticks = {}
    min_value = 0
    max_value = 0

    data = sample_wafer_data['Data']
    if sample_wafer_data['Unit'] is None:
        name = list(data.keys())[0]
        ND[name] = len(data[name]['measure'])
        ND_Ticks[name] = (0, data[name]['name'])
    else:
        all_measure = []
        for cnt, (key, value) in enumerate(data.items()):
            measures = list(map(float, value['measure'].values()))
            ND[key] = len(measures)
            ND_Ticks[key] = (cnt, value['name'])
            all_measure.extend(measures)

        min_value, max_value = min(all_measure), max(all_measure)
        space_value = int(max_value - min_value) + 1
        space = 1 if space_value < 100 else space_value // 100

        for i in range(space_value):
            bin_min = space * i + min_value
            RD[i] = sum(1 for j in all_measure if j < bin_min)
            RD_Ticks[i] = (i, str(bin_min))

    bin_data = {
        "ND": {1: ND},
        "RD": {1: RD},
        "ND_Ticks": ND_Ticks,
        "RD_Ticks": RD_Ticks,
        "min": min_value,
        "max": max_value,
    }

    return bin_data


def get_waferSummary(df, precision, test_name, signal, split: int = 5):
    sample_wafer_data = {}
    names_data = testNames_measure(df, precision, test_name, split=split)
    name_Signal_data = testNames_measure(df, precision, test_name, signal, split)
    if name_Signal_data:
        sample_wafer_data.update({
            'ND': name_Signal_data['ND'],
            'RD': name_Signal_data['RD'],
            'ND_Ticks': name_Signal_data['ND_Ticks'],
            'RD_Ticks': name_Signal_data['RD_Ticks'],
            'RDmin': name_Signal_data['min'],
            'RDmax': name_Signal_data['max'],
            'RDtestName': f"{test_name} {signal}"
        })

    if names_data:
        sample_wafer_data.update({
            'ALLND': names_data['ND'],
            'ALLRD': names_data['RD'],
            'ALLND_Ticks': names_data['ND_Ticks'],
            'ALLRD_Ticks': names_data['RD_Ticks'],
            'ALLRDmin': names_data['min'],
            'ALLRDmax': names_data['max'],
            'ALLRDtestName': test_name
        })
    return sample_wafer_data


def testNames_measure(df, precision: dict, test_name: str = None, signal: str = None,
                      split: int = 5):
    # Extract columns containing "Measure"

    measure_cols = [col for col in df.columns if 'Measure' in col]
    if not measure_cols:
        return {}

    # Filter rows by test name
    filter_condition = (df['TestName'] == test_name)
    if signal:
        filter_condition &= (df['Signal'] == signal)

    filtered_df = df[filter_condition]
    if filtered_df.empty:
        return {}

    # Unit conversion dictionary
    unit_conversion = {
        "mV": 1, "V": 1000, "uV": 0.001, "nV": 0.000001,
        "mA": 1, "A": 1000, "uA": 0.001, "nA": 0.000001,
    }

    # Extract and convert values
    all_measure = []
    units = {}
    for measure in measure_cols:
        values_with_units = filtered_df[measure].dropna()
        for value in values_with_units:
            value_and_unit = extract_value_and_unit(value)
            if value_and_unit:
                unit = value_and_unit['unit']
                value = float(value_and_unit['value'])
                pre = precision.get(unit, 3)
                converted_value = round(value * unit_conversion.get(unit, 1), pre)
                all_measure.append(converted_value)
                units[unit] = unit


    if not all_measure:
        return {}

    # Calculate min, max, and binning
    min_value = min(all_measure)
    max_value = max(all_measure)
    space_value = int(max_value - min_value) + 1
    space = 1 if space_value < 100 else space_value // 100

    # Relative Distribution (RD)
    RD = {i: sum(1 for value in all_measure if value < (space * i + min_value))
          for i in range(space_value)}
    RD_Ticks = {i: (i, str(space * i + min_value)) for i in range(space_value)}

    split_edges = np.linspace(min_value, max_value, split + 1)
    ND = {}
    ND_Ticks = {}

    for i in range(split):
        set_color = f'split_{i}'
        lower_bound = split_edges[i]
        upper_bound = split_edges[i + 1]
        min_measure = float('inf')
        max_measure = float('-inf')
        measure_cnt = 0
        # Assign each value to the appropriate bin
        for value in all_measure:
            if lower_bound <= value < upper_bound:
                measure_cnt += 1
                # Update min and max measure counts
                if max_measure < value:
                    max_measure = value
                if min_measure > value:
                    min_measure = value
            elif i == split - 1 and value == upper_bound:
                measure_cnt += 1
                # Update min and max measure counts
                if max_measure < value:
                    max_measure = value
                if min_measure > value:
                    min_measure = value
        if measure_cnt:
            ND[set_color] = measure_cnt
            ND_Ticks[set_color] = (i, f'[{min_measure}~{max_measure}]')

    # Compile results
    bin_data = {
        "ND": {1: ND},
        "RD": {1: RD},
        "ND_Ticks": ND_Ticks,
        "RD_Ticks": RD_Ticks,
        "min": min_value,
        "max": max_value,
    }
    return bin_data


def extract_value_and_unit(input_string):
    # 正则表达式匹配必须有一个数值和一个单位
    regex = r'^([+-]?\d*\.?\d+)\s*(mV|mA|uA|nA|A|uV|nV|V)$'
    match = re.match(regex, input_string)

    if match:
        # 第一个括号匹配的是数值部分
        value = float(match.group(1))
        # 第二个括号匹配的是单位部分
        unit = match.group(2)

        return {"value": value, "unit": unit}
    else:
        # 如果输入格式不匹配，则返回None或其他适当的错误处理
        return None


if __name__ == '__main__':
    df = pd.read_csv(r"D:\dataset\new3\auto_20240829-142329_datalog_19700101080000.csv", dtype=str, low_memory=False)

    # sample_wafer_data = get_waferSummary(df, 'MONI1', 'Vpph')
    sample_wafer_data = testNames_measure(df,{}, 'MONI1')
    app = QApplication(sys.argv)
    wafer_map = WaferSummary()
    wafer_map.setSummaryData(sample_wafer_data)
    wafer_map.show()
    sys.exit(app.exec_())
